from matplotlib import pyplot as plt
from matplotlib.ticker import MultipleLocator
from sklearn.metrics import r2_score
import seaborn as sns
import numpy as np

plt.rcParams['font.sans-serif'] = ['Arial Unicode Ms']

font_size = 15


def plot_loss(loss, path):
    fig = plt.figure(figsize=(10, 6))
    ax = fig.add_subplot(111)
    x = list(range(1, 1 + len(loss)))
    ax.plot(x, loss)
    plt.savefig(path)


def plot_output(real, data, date, path):
    fig = plt.figure(figsize=(14, 7))
    ax = fig.add_subplot(111)
    xminorLocator = MultipleLocator(1)
    yminorLocator = MultipleLocator(0.5)
    ax.plot(date, real, label="观测值", linewidth=2.5)
    ylim = [0, 30]
    if '春_self' in path:
        ylim = [12, 26]
    elif '夏_self' in path:
        ylim = [20, 30]
    elif '秋_self' in path:
        ylim = [8, 20]
    elif '冬_self' in path:
        ylim = [0, 12]
    elif '春' in path:
        ylim = [10, 29]
    elif '夏' in path:
        ylim = [20, 32]
    elif '秋' in path:
        ylim = [8, 30]
    elif '冬' in path:
        ylim = [-2, 16]

    patten = '[%.4f, %.4f]'
    modes = ''
    colors = ['red', 'blue']
    for key, color in zip(data, colors):
        pred = data[key]
        ax.plot(date, pred, label="%s预测" % key, linewidth=2.5)
        # ax.scatter(date, pred, color=color)
        r2 = r2_score(real, pred)
        r = np.sqrt(r2)
        eval_result = patten % (r, r2)
        modes += '%s_%s_' % (key, eval_result)

    ax.set_ylim(ylim)
    ax.set_yticks(np.arange(ylim[0], ylim[1] + 0.1, 2))

    ax.set_xlim([date[0], date[-1]])
    m = date[::6]
    m.append(date[-1])
    ax.set_xticks(m)

    ax.xaxis.set_minor_locator(xminorLocator)
    ax.yaxis.set_minor_locator(yminorLocator)
    plt.xlabel("时间/10min",  fontsize=font_size)
    plt.ylabel("气温/℃", fontsize=font_size)
    plt.legend(loc=2, fontsize=font_size)
    plt.grid(which="minor", ls='--', alpha=0.3)
    plt.grid()
    plt.savefig(path + '_%s.png' % modes, dpi=300)
    return modes

def plot_corr(df, path):
    fig = plt.figure(figsize=(12, 8))
    ax = fig.add_subplot(111)
    sns.heatmap(df.corr(), annot=True)
    plt.savefig(path + '.png', dpi=400)


def show_predict_result(real, pred, filename, xlabel="实测值(℃)", ylabel="观测值(℃)"):
    fig = plt.figure(figsize=(8, 8))
    ax = fig.add_subplot(111)
    Max = max(max(real), max(pred))
    Min = min(min(real), min(pred))
    f = np.polyfit(real, pred, 1)
    f = np.poly1d(f)

    ax.scatter(real, pred, marker='.')
    ax.set_xlim([Min, Max])
    ax.set_ylim([Min, Max])
    plt.xticks(fontsize=font_size)
    plt.yticks(fontsize=font_size)
    x = np.linspace(Min, Max, 10)
    ax.plot(x, x, color="red", label="趋势线")
    y = f(x)
    ax.plot(x, y, linestyle='--', color="red", label="1:1线")
    plt.scatter([0], [0], alpha=0, label="R²=%.4f" % r2_score(real, pred))
    plt.scatter([0], [0], alpha=0, label="n=%d" % real.shape[0])
    ax.set_xlabel(xlabel, fontsize=font_size)
    ax.set_ylabel(ylabel, fontsize=font_size)
    plt.legend(loc="upper left", fontsize=font_size)
    plt.savefig(filename, dpi=400)
